# Copyright (c) 2017 Presto Labs Pte. Ltd.
# Author: inkyu

import logging
import json
from typing import List, Union

from tornado import gen
from numpy import array_split

import coin.flow.topic_map as topic_map
from coin.base.param_util import to_list
from coin.exchange.base.ws.subscriber import WebsocketJsonSubscriberBase


class BitfinexV2Subscriber(WebsocketJsonSubscriberBase):
  def __init__(self,
               queue,
               group_id: str,
               worker_id: str,
               symbol: Union[str, List[str]],
               close_callback=None,
               logger=None):
    super().__init__(close_callback, logger or logging.getLogger(__name__))
    self._symbol = to_list(symbol)
    self._channels = {}
    self._register_url('wss://api.bitfinex.com/ws/2')

    self._publisher = None
    self._init_publisher(queue, group_id, worker_id)

  def _init_publisher(self, queue, group_id: str, worker_id: str):
    topic_book = topic_map.search_topic_id('bitfinex_v2_ws_book_%s%s' % (group_id, worker_id))[0]
    topic_raw_book = topic_map.search_topic_id('bitfinex_v2_ws_rawBook_%s%s' %
                                               (group_id, worker_id))[0]
    topic_trades = topic_map.search_topic_id('bitfinex_v2_ws_trades_%s%s' %
                                             (group_id, worker_id))[0]
    topic_etc = topic_map.search_topic_id('bitfinex_v2_ws_etc_%s%s' % (group_id, worker_id))[0]

    self._publisher = {
        'book': queue.get_writer(topic_id=topic_book[1]),
        'rawBook': queue.get_writer(topic_id=topic_raw_book[1]),
        'trades': queue.get_writer(topic_id=topic_trades[1]),
        'etc': queue.get_writer(topic_id=topic_etc[1])
    }

    self._enable_topic_event_publisher(
        queue,
        'bitfinex_v2_ws_topic_event',
        [topic_book[0], topic_raw_book[0], topic_trades[0], topic_etc[0]])

    self._add_teardown_callback(self._teardown_publisher)

  def _teardown_publisher(self):
    if self._publisher:
      for _, publisher in self._publisher.items():
        publisher.close()
      self._publisher = None

  @gen.coroutine
  def on_open(self):
    msg = yield self._ws.read_message()
    msg = json.loads(msg)
    if msg.get('event', None) != 'info' or msg.get('version', None) != 2:
      self._logger.error('Version mismatch: "%s". Stop subscription.' % msg)
      self._notify_close()
      return
    self._ioloop.add_callback(self._send_subscribe_request)

  @gen.coroutine
  def _send_subscribe_request(self):
    for symbol in self._symbol:
      self._logger.debug('Subscribing %s' % symbol)
      yield self._ws.write_message(
          json.dumps({
              'event': 'subscribe', 'channel': 'book', 'symbol': symbol, 'len': '100'
          }))
      yield self._ws.write_message(
          json.dumps({
              'event': 'subscribe', 'channel': 'trades', 'symbol': symbol
          }))
      yield self._ws.write_message(
          json.dumps({
              'event': 'subscribe', 'channel': 'book', 'symbol': symbol, 'prec': 'R0', 'len': '100'
          }))
      yield gen.sleep(0.5)  # Sleep 0.5 second before next request

    self._publish_topic_event('START')
    self._logger.info('Subscribe done %d', id(self))

  def on_message(self, timestamp: int, message: Union[dict, list]):
    if self.stopped:
      return

    topic = 'etc'
    if type(message) == list:
      channel = message[0]
      try:
        topic = self._channels[channel]
      except KeyError:
        self._logger.warning('Unknown channel: %s. Known channels are %s' %
                             (channel, self._channels))
        return
    elif type(message) == dict:
      if message['event'] == 'subscribed':
        topic = message['channel']
        channel = message['chanId']
        if topic == 'book' and message['prec'] == 'R0':
          topic = 'rawBook'
        self._channels[channel] = topic

    try:
      self._publisher[topic].write(timestamp, message)
    except KeyError:
      self._logger.error('Unknown topic: %s' % topic)


# Wrapper class splits symbols into no more than 70 for each group.
# This is because Bitfinex have subscription limit of 250 and each symbol
# subscribes book, trades and raw book.
# https://docs.bitfinex.com/v2/docs/ws-general
# Subscription Limit
# Starting September 15th, 2018, 00:00 UTC All websocket connections will have
# a limit of 250 subscriptions to market data feed channels (tickers, book,
# candles, trades, …). Users using or requiring more than 250 subscriptions
# will need to split their subscriptions to channels using multiple WebSocket
# Connections.
class BitfinexV2SubscriberWrapper(object):
  def __init__(self,
               queue,
               group_id: str,
               worker_id: str,
               symbols: Union[List[str], str],
               close_callback=None,
               logger=None):
    self._subscribers = []
    for small_group in array_split(symbols, len(symbols) / 70 + 1):
      self._subscribers.append(
          BitfinexV2Subscriber(queue,
                               group_id,
                               worker_id,
                               small_group.tolist(),
                               close_callback,
                               logger))

  def start(self, *args, **kwargs):
    for subscriber in self._subscribers:
      subscriber.start(*args, **kwargs, delay_time=30)

  def stop(self, *args, **kwargs):
    for subscriber in self._subscribers:
      subscriber.stop(*args, **kwargs)
